// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <gtest/gtest.h>

#include <pollux/common/memory/raw_vector.h>
#include <pollux/testing/vector/vector_test_base.h>

using namespace kumo::pollux;
using namespace kumo::pollux::test;

class LazyVectorTest : public testing::Test, public VectorBuilder {
protected:
    static void SetUpTestCase() {
        memory::MemoryManager::testingSetInstance({});
    }
};

TEST_F(LazyVectorTest, lazyInDictionary) {
    // We have a dictionary over LazyVector. We load for some indices in
    // the dictionary. We check that the loads on the wrapped lazy
    // vector are properly translated and deduplicated.
    static constexpr int32_t kInnerSize = 100;
    static constexpr int32_t kOuterSize = 1000;
    std::vector<vector_size_t> loadedRows;
    auto lazy = std::make_shared<LazyVector>(
        pool_.get(),
        INTEGER(),
        kInnerSize,
        std::make_unique<SimpleVectorLoader>([&](auto rows) {
            for (auto row: rows) {
                loadedRows.push_back(row);
            }
            return make_flat_vector<int32_t>(
                rows.back() + 1, [](auto row) { return row; });
        }));
    auto wrapped = BaseVector::wrap_in_dictionary(
        nullptr,
        make_indices(kOuterSize, [](auto row) { return row / 10; }),
        kOuterSize,
        lazy);

    // We expect a single level of dictionary and rows loaded for the selected
    // indices in rows.

    SelectivityVector rows(kOuterSize, false);
    // We select 3 rows, the 2 first fall on 0 and the last on 5 in 'base'.
    rows.setValid(1, true);
    rows.setValid(9, true);
    rows.setValid(55, true);
    rows.updateBounds();
    LazyVector::ensureLoadedRows(wrapped, rows);
    EXPECT_EQ(wrapped->encoding(), VectorEncoding::Simple::DICTIONARY);
    EXPECT_EQ(wrapped->value_vector()->encoding(), VectorEncoding::Simple::FLAT);

    EXPECT_EQ(loadedRows, (std::vector<vector_size_t>{0, 5}));
    kumo::pollux::test::assertCopyableVector(wrapped);
}

// Ensure that compare() using lazies and lazies wrapped by a dictionay work
// properly.
TEST_F(LazyVectorTest, compareLazies) {
    static constexpr int32_t kInnerSize = 100;
    static constexpr int32_t kOuterSize = 100;

    auto getLazy = [&]() {
        return std::make_shared<LazyVector>(
            pool_.get(),
            INTEGER(),
            kInnerSize,
            std::make_unique<SimpleVectorLoader>([&](auto) {
                return make_flat_vector<int32_t>(
                    kInnerSize, [](auto row) { return row; });
            }));
    };

    auto getDictionary = [&]() {
        return BaseVector::wrap_in_dictionary(
            nullptr,
            make_indices(kOuterSize, [](auto row) { return row; }),
            kOuterSize,
            getLazy());
    };

    auto expected =
            make_flat_vector<int32_t>(kInnerSize, [](auto row) { return row; });

    // First compare with lazies.
    auto lazy1 = getLazy();
    auto lazy2 = getLazy();

    for (size_t i = 0; i < kInnerSize; ++i) {
        // Compare with lazies on either side.
        EXPECT_EQ(lazy1->compare(expected.get(), i, i, {}), 0);
        EXPECT_EQ(expected->compare(lazy2.get(), i, i, {}), 0);
    }

    // Then compare with lazies inside a dictionary.
    auto dictionaryLazy1 = getDictionary();
    auto dictionaryLazy2 = getDictionary();

    for (size_t i = 0; i < kOuterSize; ++i) {
        // Compare with dictionaries wrapped around lazies on either side.
        EXPECT_EQ(dictionaryLazy1->compare(expected.get(), i, i, {}), 0);
        EXPECT_EQ(expected->compare(dictionaryLazy2.get(), i, i, {}), 0);
    }
}

TEST_F(LazyVectorTest, rowVectorWithLazyChild) {
    constexpr vector_size_t size = 1000;
    auto columnType = ROW({"a", "b"}, {INTEGER(), INTEGER()});

    auto lazyVectorA = vectorMaker_.lazy_flat_vector<int32_t>(
        size,
        [&](vector_size_t i) { return i % 5; },
        [](vector_size_t i) { return i % 7 == 0; });
    auto lazyVectorB = vectorMaker_.lazy_flat_vector<int32_t>(
        size,
        [&](vector_size_t i) { return i % 3; },
        [](vector_size_t i) { return i % 11 == 0; });

    VectorPtr row_vector = make_row_vector({lazyVectorA, lazyVectorB});
    EXPECT_TRUE(is_lazy_not_loaded(*row_vector.get()));

    SelectivityVector rows(row_vector->size(), false);
    LazyVector::ensureLoadedRows(row_vector, rows);
    EXPECT_FALSE(is_lazy_not_loaded(*row_vector.get()));
}

TEST_F(LazyVectorTest, rowVectorWithLazyChildPartiallyLoaded) {
    constexpr vector_size_t size = 1000;
    auto columnType = ROW({"a", "b"}, {INTEGER(), INTEGER()});

    auto lazyVectorA = vectorMaker_.lazy_flat_vector<int32_t>(
        size, [&](vector_size_t) { return 2222; });
    auto lazyVectorB = vectorMaker_.lazy_flat_vector<int32_t>(
        size, [&](vector_size_t) { return 1111; });

    VectorPtr row_vector = make_row_vector({lazyVectorA, lazyVectorB});
    EXPECT_FALSE(lazyVectorA->isLoaded());
    EXPECT_TRUE(is_lazy_not_loaded(*row_vector.get()));

    // Only load first row, the row vector should still be valid.
    SelectivityVector rows(1);
    LazyVector::ensureLoadedRows(row_vector, rows);
    EXPECT_FALSE(is_lazy_not_loaded(*row_vector.get()));

    auto rowVectorPtr = row_vector->as_unchecked<RowVector>();
    auto child1 = rowVectorPtr->childAt(0);
    auto child2 = rowVectorPtr->childAt(1);

    EXPECT_FALSE(child1->is_null_at(0));
    EXPECT_FALSE(child2->is_null_at(0));
    EXPECT_EQ(child1->loaded_vector()->as_flat_vector<int32_t>()->value_at(0), 2222);
    EXPECT_EQ(child2->loaded_vector()->as_flat_vector<int32_t>()->value_at(0), 1111);
    EXPECT_FALSE(rowVectorPtr->childAt(1)->is_null_at(0));

    EXPECT_EQ(rowVectorPtr->childAt(0)->size(), 1000);
    EXPECT_EQ(rowVectorPtr->childAt(1)->size(), 1000);
}

TEST_F(LazyVectorTest, dictionaryOverRowVectorWithLazyChild) {
    constexpr vector_size_t size = 1000;
    auto columnType = ROW({"a", "b"}, {INTEGER(), INTEGER()});

    auto lazyVectorA = vectorMaker_.lazy_flat_vector<int32_t>(
        size,
        [&](vector_size_t i) { return i % 5; },
        [](vector_size_t i) { return i % 7 == 0; });
    auto lazyVectorB = vectorMaker_.lazy_flat_vector<int32_t>(
        size,
        [&](vector_size_t i) { return i % 3; },
        [](vector_size_t i) { return i % 11 == 0; });

    VectorPtr row_vector = make_row_vector({lazyVectorA, lazyVectorB});
    VectorPtr dict = BaseVector::wrap_in_dictionary(
        nullptr,
        make_indices(size, [](auto row) { return row; }),
        size,
        BaseVector::wrap_in_dictionary(
            nullptr,
            make_indices(size, [](auto row) { return row; }),
            size,
            row_vector));
    EXPECT_TRUE(is_lazy_not_loaded(*dict.get()));

    SelectivityVector rows(dict->size(), false);
    LazyVector::ensureLoadedRows(dict, rows);
    EXPECT_FALSE(is_lazy_not_loaded(*dict.get()));
    // Ensure encoding layer is correctly initialized after lazy loading.
    EXPECT_NO_THROW(dict->may_have_nulls_recursive());
}

TEST_F(LazyVectorTest, selectiveRowVectorWithLazyChild) {
    constexpr vector_size_t size = 1000;
    auto columnType = ROW({"a", "b"}, {INTEGER(), INTEGER()});
    int loadedA = 0, loadedB = 0;
    int expectedLoadedA = 0, expectedLoadedB = 0;

    auto lazyVectorA =
            vectorMaker_.lazy_flat_vector<int32_t>(size, [&](vector_size_t i) {
                ++loadedA;
                return i % 5;
            });
    auto lazyVectorB =
            vectorMaker_.lazy_flat_vector<int32_t>(size, [&](vector_size_t i) {
                ++loadedB;
                return i % 3;
            });

    VectorPtr row_vector = make_row_vector({lazyVectorA, lazyVectorB});
    EXPECT_TRUE(is_lazy_not_loaded(*row_vector.get()));

    SelectivityVector rows(row_vector->size(), false);
    for (int i = 0; i < size; ++i) {
        if (i % 7) {
            rows.setValid(i, true);
            ++expectedLoadedA;
            ++expectedLoadedB;
        }
    }
    rows.updateBounds();
    LazyVector::ensureLoadedRows(row_vector, rows);
    EXPECT_FALSE(is_lazy_not_loaded(*row_vector.get()));
    EXPECT_LT(expectedLoadedA, size);
    EXPECT_EQ(loadedA, expectedLoadedA);
    EXPECT_EQ(loadedB, expectedLoadedB);
}

TEST_F(LazyVectorTest, lazyRowVectorWithLazyChildren) {
    constexpr vector_size_t size = 1000;
    auto columnType =
            ROW({"inner_row"}, {ROW({"a", "b"}, {INTEGER(), INTEGER()})});

    auto lazyVectorA = vectorMaker_.lazy_flat_vector<int32_t>(
        size,
        [](vector_size_t i) { return i % 5; },
        [](vector_size_t i) { return i % 7 == 0; });
    auto lazyVectorB = vectorMaker_.lazy_flat_vector<int32_t>(
        size,
        [](vector_size_t i) { return i % 3; },
        [](vector_size_t i) { return i % 11 == 0; });

    VectorPtr lazyInnerRowVector = std::make_shared<LazyVector>(
        pool_.get(),
        columnType->childAt(0),
        size,
        std::make_unique<SimpleVectorLoader>([&](RowSet /* rowSet */) {
            return vectorMaker_.row_vector({lazyVectorA, lazyVectorB});
        }));

    VectorPtr row_vector = std::make_shared<LazyVector>(
        pool_.get(),
        columnType,
        size,
        std::make_unique<SimpleVectorLoader>([&](RowSet /* rowSet */) {
            return vectorMaker_.row_vector({lazyInnerRowVector});
        }));

    EXPECT_TRUE(is_lazy_not_loaded(*row_vector));

    DecodedVector decodedVector;
    SelectivityVector baseRows;
    LazyVector::ensureLoadedRows(
        row_vector, SelectivityVector(size, true), decodedVector, baseRows);

    EXPECT_FALSE(is_lazy_not_loaded(*row_vector.get()));

    RowVector *loaded_vector = row_vector->loaded_vector()->as<RowVector>();
    auto child = loaded_vector->childAt(0);
    // EXPECT_FALSE(child->is_lazy());
    assertEqualVectors(child, lazyInnerRowVector);
}

TEST_F(LazyVectorTest, dictionaryOverLazyRowVectorWithLazyChildren) {
    constexpr vector_size_t size = 1000;
    auto columnType =
            ROW({"inner_row"}, {ROW({"a", "b"}, {INTEGER(), INTEGER()})});

    auto lazyVectorA = vectorMaker_.lazy_flat_vector<int32_t>(
        size,
        [](vector_size_t i) { return i % 5; },
        [](vector_size_t i) { return i % 7 == 0; });
    auto lazyVectorB = vectorMaker_.lazy_flat_vector<int32_t>(
        size,
        [](vector_size_t i) { return i % 3; },
        [](vector_size_t i) { return i % 11 == 0; });

    VectorPtr lazyInnerRowVector = std::make_shared<LazyVector>(
        pool_.get(),
        columnType->childAt(0),
        size,
        std::make_unique<SimpleVectorLoader>([&](RowSet /* rowSet */) {
            return vectorMaker_.row_vector({lazyVectorA, lazyVectorB});
        }));

    VectorPtr row_vector = std::make_shared<LazyVector>(
        pool_.get(),
        columnType,
        size,
        std::make_unique<SimpleVectorLoader>([&](RowSet /* rowSet */) {
            return vectorMaker_.row_vector({lazyInnerRowVector});
        }));

    VectorPtr dict = BaseVector::wrap_in_dictionary(
        nullptr,
        make_indices(size, [](auto row) { return row; }),
        size,
        BaseVector::wrap_in_dictionary(
            nullptr,
            make_indices(size, [](auto row) { return row; }),
            size,
            row_vector));

    EXPECT_TRUE(is_lazy_not_loaded(*dict));

    DecodedVector decodedVector;
    SelectivityVector baseRows;
    LazyVector::ensureLoadedRows(
        dict, SelectivityVector(size, true), decodedVector, baseRows);

    EXPECT_FALSE(is_lazy_not_loaded(*dict.get()));

    RowVector *loaded_vector =
            dict->value_vector()->value_vector()->loaded_vector()->as<RowVector>();
    auto child = loaded_vector->childAt(0);
    EXPECT_FALSE(child->is_lazy());
    assertEqualVectors(child, lazyInnerRowVector);
}

TEST_F(
    LazyVectorTest,
    dictionaryOverLazyRowVectorWithDictionaryOverLazyChildren) {
    // Input: dict(row(lazy(dict(row(lazy)))))
    constexpr vector_size_t size = 1000;

    auto lazyVectorA = vectorMaker_.lazy_flat_vector<int32_t>(
        size,
        [](vector_size_t i) { return i % 5; },
        [](vector_size_t i) { return i % 7 == 0; });

    VectorPtr innerRowVector = vectorMaker_.row_vector({lazyVectorA});
    VectorPtr innerDict = BaseVector::wrap_in_dictionary(
        nullptr,
        make_indices(size, [](auto row) { return row; }),
        size,
        innerRowVector);

    VectorPtr row_vector = vectorMaker_.row_vector({
        std::make_shared<LazyVector>(
            pool_.get(),
            innerRowVector->type(),
            size,
            std::make_unique<SimpleVectorLoader>(
                [&](RowSet /* rowSet */) { return innerDict; }))
    });

    VectorPtr dict = BaseVector::wrap_in_dictionary(
        nullptr,
        make_indices(size, [](auto row) { return row; }),
        size,
        row_vector);

    EXPECT_TRUE(is_lazy_not_loaded(*dict));

    DecodedVector decodedVector;
    SelectivityVector baseRows;
    LazyVector::ensureLoadedRows(
        dict, SelectivityVector(size, true), decodedVector, baseRows);

    EXPECT_FALSE(is_lazy_not_loaded(*dict.get()));

    // Ensure encoding layer is correctly initialized after lazy loading.
    EXPECT_NO_THROW(dict->may_have_nulls_recursive());
    EXPECT_NO_THROW(innerDict->may_have_nulls_recursive());
    // Verify that the correct vector is loaded.
    ASSERT_EQ(
        innerDict.get(),
        dict->value_vector()
        ->as_unchecked<RowVector>()
        ->childAt(0)
        ->loaded_vector());
}

TEST_F(LazyVectorTest, lazyRowVectorWithLazyChildrenPartiallyLoaded) {
    constexpr vector_size_t size = 1000;
    auto columnType =
            ROW({"inner_row"}, {ROW({"a", "b"}, {INTEGER(), INTEGER()})});

    auto lazyVectorA = vectorMaker_.lazy_flat_vector<int32_t>(
        size,
        [](vector_size_t i) { return i % 5; },
        [](vector_size_t i) { return i % 7 == 0; });
    auto lazyVectorB = vectorMaker_.lazy_flat_vector<int32_t>(
        size,
        [](vector_size_t i) { return i % 3; },
        [](vector_size_t i) { return i % 11 == 0; });

    VectorPtr lazyInnerRowVector = std::make_shared<LazyVector>(
        pool_.get(),
        columnType->childAt(0),
        size,
        std::make_unique<SimpleVectorLoader>([&](RowSet /* rowSet */) {
            return vectorMaker_.row_vector({lazyVectorA, lazyVectorB});
        }));

    VectorPtr row_vector = std::make_shared<LazyVector>(
        pool_.get(),
        columnType,
        size,
        std::make_unique<SimpleVectorLoader>([&](RowSet /* rowSet */) {
            return vectorMaker_.row_vector({lazyInnerRowVector});
        }));

    EXPECT_TRUE(is_lazy_not_loaded(*row_vector));

    raw_vector<vector_size_t> rowNumbers;
    auto iota = kumo::pollux::iota(size, rowNumbers);
    RowSet rowSet(iota, size);
    row_vector->as_unchecked<LazyVector>()->load(rowSet, nullptr);
    EXPECT_TRUE(is_lazy_not_loaded(*row_vector));

    DecodedVector decodedVector;
    SelectivityVector baseRows;
    LazyVector::ensureLoadedRows(
        row_vector, SelectivityVector(size, true), decodedVector, baseRows);

    EXPECT_FALSE(is_lazy_not_loaded(*row_vector.get()));

    RowVector *loaded_vector = row_vector->loaded_vector()->as<RowVector>();
    auto child = loaded_vector->childAt(0);
    EXPECT_FALSE(is_lazy_not_loaded(*child));
    assertEqualVectors(child, lazyInnerRowVector);
}

TEST_F(LazyVectorTest, lazyInConstant) {
    // Wrap Lazy vector in a Constant, load some indices and verify that the
    // results.
    static constexpr int32_t kInnerSize = 100;
    static constexpr int32_t kOuterSize = 1000;
    auto base = make_flat_vector<int32_t>(kInnerSize, [](auto row) { return row; });
    std::vector<vector_size_t> loadedRows;
    auto lazy = std::make_shared<LazyVector>(
        pool_.get(),
        INTEGER(),
        kInnerSize,
        std::make_unique<SimpleVectorLoader>([&](auto rows) {
            for (auto row: rows) {
                loadedRows.push_back(row);
            }
            return base;
        }));
    VectorPtr wrapped =
            std::make_shared<ConstantVector<int32_t> >(pool_.get(), 10, 7, lazy);

    SelectivityVector rows(kOuterSize, false);
    rows.setValid(1, true);
    rows.setValid(9, true);
    rows.setValid(55, true);
    rows.updateBounds();
    LazyVector::ensureLoadedRows(wrapped, rows);
    EXPECT_EQ(wrapped->encoding(), VectorEncoding::Simple::CONSTANT);
    EXPECT_EQ(loadedRows, (std::vector<vector_size_t>{7}));

    EXPECT_EQ(wrapped->as<SimpleVector<int32_t>>()->value_at(1), 7);
    EXPECT_EQ(wrapped->as<SimpleVector<int32_t>>()->value_at(9), 7);
    EXPECT_EQ(wrapped->as<SimpleVector<int32_t>>()->value_at(55), 7);
}

TEST_F(LazyVectorTest, lazyInDoubleDictionary) {
    // We have dictionaries over LazyVector. We load for some indices in
    // the top dictionary. The intermediate dictionaries refer to
    // non-loaded items in the base of the LazyVector, including indices
    // past its end.
    // This test make sure the we have valid vector after loading.
    static constexpr int32_t kInnerSize = 100;
    static constexpr int32_t kOuterSize = 1000;

    VectorPtr lazy;
    vector_size_t loadEnd = 0;

    auto makeWrapped = [&](BufferPtr nulls) {
        loadEnd = 0;
        lazy = std::make_shared<LazyVector>(
            pool_.get(),
            INTEGER(),
            kOuterSize,
            std::make_unique<SimpleVectorLoader>([&](auto rows) {
                loadEnd = rows.back() + 1;
                return make_flat_vector<int32_t>(loadEnd, [](auto row) { return row; });
            }));

        return BaseVector::wrap_in_dictionary(
            std::move(nulls),
            make_indices(kInnerSize, [](auto row) { return row; }),
            kInnerSize,
            BaseVector::wrap_in_dictionary(
                nullptr,
                make_indices(kOuterSize, [](auto row) { return row; }),
                kOuterSize,
                lazy));
    };
    SelectivityVector rows(kInnerSize);

    // No nulls.
    {
        auto wrapped = makeWrapped(nullptr);

        LazyVector::ensureLoadedRows(wrapped, rows);
        EXPECT_EQ(kInnerSize, loadEnd);
        auto expected =
                make_flat_vector<int32_t>(kInnerSize, [](auto row) { return row; });
        assertEqualVectors(wrapped, expected);
    }

    // With nulls.
    {
        auto wrapped = makeWrapped(make_nulls(kInnerSize, null_every(7)));
        LazyVector::ensureLoadedRows(wrapped, rows);

        EXPECT_EQ(kInnerSize, loadEnd);
        auto expected = make_flat_vector<int32_t>(
            kInnerSize, [](auto row) { return row; }, null_every(7));
        assertEqualVectors(wrapped, expected);
    }

    // With nulls at the end.
    {
        auto wrapped = makeWrapped(make_nulls(kInnerSize, null_every(3)));
        LazyVector::ensureLoadedRows(wrapped, rows);

        EXPECT_EQ(kInnerSize - 1, loadEnd);

        auto expected = make_flat_vector<int32_t>(
            kInnerSize, [](auto row) { return row; }, null_every(3));
        assertEqualVectors(wrapped, expected);
    }
}

TEST_F(LazyVectorTest, lazySlice) {
    auto lazy = std::make_shared<LazyVector>(
        pool_.get(),
        INTEGER(),
        100,
        std::make_unique<SimpleVectorLoader>([&](auto rows) {
            return make_flat_vector<int32_t>(
                rows.back() + 1, [](auto row) { return row; });
        }));
    EXPECT_THROW(lazy->slice(0, 10), PolluxRuntimeError);
    lazy->loaded_vector();
    auto slice = lazy->slice(0, 10);
    for (int i = 0; i < slice->size(); ++i) {
        EXPECT_TRUE(slice->equal_value_at(lazy.get(), i, i));
    }
}

TEST_F(LazyVectorTest, lazyInMultipleDictionaryAllResultantNullRows) {
    // Verifies that lazy loading works for a lazy vector that is wrapped in
    // multiple layers of dictionary encoding such that the rows that it needs to
    // be loaded for all end up pointing to nulls. This results in a zero sized
    // base vector which when wrapped in a dictionary layer can run into invalid
    // internal state for row indices that were not asked to be loaded.
    static constexpr int32_t kVectorSize = 10;
    auto lazy = std::make_shared<LazyVector>(
        pool_.get(),
        INTEGER(),
        kVectorSize,
        std::make_unique<SimpleVectorLoader>([&](auto rows) {
            return make_flat_vector<int32_t>(
                rows.back() + 1, [](auto row) { return row; });
        }));
    auto wrapped = BaseVector::wrap_in_dictionary(
        make_nulls(kVectorSize, [](vector_size_t /*row*/) { return true; }),
        make_indices(kVectorSize, [](auto row) { return row; }),
        kVectorSize,
        lazy);
    wrapped = BaseVector::wrap_in_dictionary(
        nullptr,
        make_indices(kVectorSize, [](auto row) { return row; }),
        kVectorSize,
        wrapped);
    SelectivityVector rows(kVectorSize, true);
    rows.setValid(1, false);
    LazyVector::ensureLoadedRows(wrapped, rows);
    auto expected =
            BaseVector::create_null_constant(lazy->type(), wrapped->size(), pool());
    assertEqualVectors(expected, wrapped);
}

TEST_F(LazyVectorTest, lazyInDictionaryNoRowsToLoad) {
    // Verifies that lazy loading works for a lazy vector that is wrapped a
    // dictionary with no extra nulls when loading for 0 selected rows.
    static constexpr int32_t kVectorSize = 10;
    auto lazy = std::make_shared<LazyVector>(
        pool_.get(),
        INTEGER(),
        kVectorSize,
        std::make_unique<SimpleVectorLoader>([&](auto rows) {
            return make_flat_vector<int32_t>(
                rows.back() + 1, [](auto row) { return row; });
        }));
    auto wrapped = BaseVector::wrap_in_dictionary(
        nullptr,
        make_indices(kVectorSize, [](auto row) { return row; }),
        kVectorSize,
        lazy);
    SelectivityVector rows(kVectorSize, false);
    LazyVector::ensureLoadedRows(wrapped, rows);
    for (int i = 0; i < wrapped->size(); i++) {
        EXPECT_TRUE(wrapped->is_null_at(i));
    }
}

TEST_F(LazyVectorTest, lazyWithDictionaryInConstant) {
    // Wrap Lazy vector in a Dictionary in a Lazy vector in a constant, load some
    // indices and verify that the results.
    static constexpr int32_t kInnerSize = 1000;
    static constexpr int32_t kOuterSize = 10;
    auto base = make_flat_vector<int32_t>(kInnerSize, [](auto row) { return row; });
    auto innerLazy = std::make_shared<LazyVector>(
        pool_.get(),
        INTEGER(),
        kInnerSize,
        std::make_unique<SimpleVectorLoader>(
            [&](auto /* rows */) { return base; }));
    auto wrapped = BaseVector::wrap_in_dictionary(
        nullptr,
        make_indices(kInnerSize, [](auto row) { return row; }),
        kInnerSize,
        innerLazy);
    auto outerLazy = std::make_shared<LazyVector>(
        pool_.get(),
        INTEGER(),
        kInnerSize,
        std::make_unique<SimpleVectorLoader>(
            [&](auto /* rows */) { return wrapped; }));
    VectorPtr constant = std::make_shared<ConstantVector<int32_t> >(
        pool_.get(), kOuterSize, 7, outerLazy);

    SelectivityVector rows(kOuterSize, true);
    LazyVector::ensureLoadedRows(constant, rows);
    EXPECT_EQ(constant->encoding(), VectorEncoding::Simple::CONSTANT);

    for (int i = 0; i < kOuterSize; i++) {
        EXPECT_EQ(constant->as<SimpleVector<int32_t>>()->value_at(i), 7);
    }
}

TEST_F(LazyVectorTest, reset) {
    static constexpr int32_t kVectorSize = 10;
    auto loader = [&](RowSet rows) {
        return make_flat_vector<int32_t>(
            rows.back() + 1, [](auto row) { return row; });
    };
    LazyVector lazy(
        pool_.get(),
        INTEGER(),
        kVectorSize,
        std::make_unique<SimpleVectorLoader>(loader));
    lazy.set_null(0, true);
    lazy.reset(std::make_unique<SimpleVectorLoader>(loader), kVectorSize);
    ASSERT_EQ(lazy.nulls(), nullptr);
}

TEST_F(LazyVectorTest, runtimeStats) {
    TestRuntimeStatWriter writer;
    RuntimeStatWriterScopeGuard guard(&writer);
    auto lazy = std::make_shared<LazyVector>(
        pool_.get(),
        INTEGER(),
        10,
        std::make_unique<SimpleVectorLoader>([&](auto rows) {
            return make_flat_vector<int32_t>(rows.back() + 1, melon::identity);
        }));
    ASSERT_EQ(lazy->loaded_vector()->size(), 10);
    auto stats = writer.stats();
    std::sort(stats.begin(), stats.end(), [](auto &x, auto &y) {
        return x.first < y.first;
    });
    ASSERT_EQ(stats.size(), 3);
    ASSERT_EQ(stats[0].first, LazyVector::kCpuNanos);
    ASSERT_GE(stats[0].second.value, 0);
    ASSERT_EQ(stats[1].first, LazyVector::kInputBytes);
    ASSERT_GE(stats[1].second.value, 0);
    ASSERT_EQ(stats[2].first, LazyVector::kWallNanos);
    ASSERT_GE(stats[2].second.value, 0);
}
